import pickle
import new
import StringIO
import pattern.crossref


def code(*args, **kwargs):
    return new.code(*args, **kwargs)

def function(*args, **kwargs):
    return new.function(*args, **kwargs)


class PluggablePickler(pickle.Pickler):

    def __init__(self, *args):
        pickle.Pickler.__init__(self, *args)
        self.dispatch = self.dispatch.copy()

    @classmethod
    def dumps(cls, obj):
        buf = StringIO.StringIO()
        cls(buf).dump(obj)
        return buf.getvalue()


class PicklerWithCode(PluggablePickler):

    def __init__(self, *args):
        PluggablePickler.__init__(self, *args)
        self.dispatch[new.code] = self.save_code.im_func
        self.dispatch[new.function] = self.save_function.im_func
        self.functions_to_fully = set()

    def save_code(self, codeobj):
        self.save_reduce(code, self._code_to_tuple(codeobj), obj=codeobj)

    def save_function(self, funcobj):
        if funcobj in self.functions_to_fully:
            self.save_reduce(function, (funcobj.func_code,{}), obj=funcobj)
        else:
            self.save_global(funcobj)

    def _code_to_tuple(self, code):
        ctor_order = ["co_argcount", "co_nlocals", "co_stacksize",
                      "co_flags", "co_code", "co_consts", "co_names",
                      "co_varnames", "co_filename", "co_name", 
                      "co_firstlineno", "co_lnotab", "co_freevars", "co_cellvars"]

        return tuple([getattr(code, x) for x in ctor_order])

    

class PicklerWithRestrictedXRef(PluggablePickler):

    def __init__(self, *args):
        PluggablePickler.__init__(self, *args)
        self.dispatch[pattern.crossref.XRef] = self.save_xref.im_func
        self.allowed_xrefs = set()

    def save_xref(self, xref):
        if xref() in self.allowed_xrefs:
            self.save_reduce(obj=xref, *xref.__reduce__())
        else:
            self.save_none(None)


class ShapesPickler(PicklerWithRestrictedXRef):

    def dump(self, shapes):
        """
        @param shapes either a list of shapes, or a tuple of such lists
        """
        if isinstance(shapes, tuple):
            for x in shapes:
                self.allowed_xrefs.update(x)
        else:
            self.allowed_xrefs.update(shapes)
        PicklerWithRestrictedXRef.dump(self, shapes)

    @staticmethod
    def cleanup(shapes):
        for shape in shapes:
            if None in shape.xref:
                del shape.xref[None]
